import os
import sys
import traceback

import numpy as np
import torch
import onnxruntime as ort

# -----------------------------
# 1. 固定随机种子 & 准备输入
# -----------------------------
np.random.seed(42)
torch.manual_seed(42)

# YOLOv7 标准输入尺寸
IMG_SIZE = 640
input_tensor = torch.randn(1, 3, IMG_SIZE, IMG_SIZE, dtype=torch.float32)
input_numpy = input_tensor.numpy()

print(f"✅ 输入张量 shape: {input_numpy.shape}, dtype: {input_numpy.dtype}")

# -----------------------------
# 2. PyTorch (.pt) 推理 (CPU)
# -----------------------------
def infer_pytorch(pt_path):
    from models.experimental import attempt_load  # 需 yolov7 代码目录在 PYTHONPATH
    model = attempt_load(pt_path, map_location='cpu')
    model.eval()
    with torch.no_grad():
        pred = model(input_tensor)[0]  # YOLOv7 输出为 list，取第一个（检测头输出）
    return pred.numpy()

# -----------------------------
# 3. ONNX Runtime (.onnx) 推理 (CPU)
# -----------------------------
def infer_onnx(onnx_path):
    sess = ort.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])
    input_name = sess.get_inputs()[0].name
    output_names = [out.name for out in sess.get_outputs()]
    outputs = sess.run(output_names, {input_name: input_numpy})
    # YOLOv7 ONNX 通常输出单个 tensor (batch, num_boxes, 85)
    return outputs[0]

# -----------------------------
# 4. 昇腾 ACL (.om) 推理 (需 CANN 环境)
# -----------------------------
def infer_om(om_path):
    """
    使用华为 ACL 接口加载 .om 模型并推理。
    注意：需在昇腾服务器上运行，并已 source env。
    """
    try:
        from infer_om import OMInference
        om_infer = OMInference()
        om_infer.load_model(om_path)
        # b8_input_numpy = np.repeat(input_numpy, 8, axis=0)  # 沿 batch 维度重复 8 次, 因为om输入为8*3*640*640
        b8_input_numpy = input_numpy
        output = om_infer.run(b8_input_numpy)
        output = output[0:1] # to 1*3*640*640
        om_infer.destroy()
    except:
        print(traceback.format_exc())
        return  None
    return output

# -----------------------------
# 6. 精度对比函数
# -----------------------------
def compare_outputs(name1, out1, name2, out2, atol=1e-4, rtol=1e-5):
    """
       在 shape 完全一致的前提下，安全地对比两个 YOLOv7 输出。

       参数:
           out1, out2: numpy.ndarray，shape 必须完全相同（如 [1, 25200, 85]）
           atol: 绝对误差容忍度（默认 1e-4）
           rtol: 相对误差容忍度（默认 1e-5）
       """
    if out1 is None or out2 is None:
        print(f"❌ 跳过 {name1} vs {name2}（缺少输出）")
        return False

    # 1. 严格检查 shape 是否一致
    if out1.shape != out2.shape:
        print(f"⚠️ Shape 不一致！{name1}: {out1.shape}, {name2}: {out2.shape}")
        return False

    # 2. 计算误差指标
    abs_diff = np.abs(out1 - out2)
    max_abs_err = np.max(abs_diff)
    mean_abs_err = np.mean(abs_diff)

    # 余弦相似度（衡量整体方向一致性）
    cosine_sim = np.dot(out1.flatten(), out2.flatten()) / (
            np.linalg.norm(out1) * np.linalg.norm(out2)
    )

    # 3. 判断是否在容忍范围内
    within_tol = np.allclose(out1, out2, atol=atol, rtol=rtol)

    # 4. 打印结果
    print(f"\n📊 {name1} vs {name2}")
    print(f"   Shape: {out1.shape}")
    print(f"   Max Abs Error: {max_abs_err:.2e}")
    print(f"   Mean Abs Error: {mean_abs_err:.2e}")
    print(f"   Cosine Similarity: {cosine_sim:.8f}")
    print(f"   AllClose (atol={atol}, rtol={rtol}): {within_tol}")

    if within_tol and cosine_sim > 0.9999:
        print("   ✅ 数值一致（差异在浮点误差范围内）")
        return True
    else:
        print("   ⚠️ 存在显著数值差异！") # 最大误差位置: (0, 23185, 2), 值: pt=393.887421, other=393.892151
        # 可选：打印最大差异位置
        idx = np.unravel_index(np.argmax(abs_diff), abs_diff.shape)
        print(f"   最大误差位置: {idx}, 值: pt={out1[idx]:.6f}, other={out2[idx]:.6f}")
        return False

# -----------------------------
# 7. 主流程
# -----------------------------
if __name__ == "__main__":
    PT_PATH = r"D:\DATA\20250814EDIPLATE\trainV7_dumpTruck\models\exp8\weights\yolov7-dump-truck-20251111060914.pt"
    ONNX_PATH = r"D:\DATA\20250814EDIPLATE\trainV7_dumpTruck\models\exp8\weights\yolov7-dump-truck-20251111060914.onnx"
    OM_PATH = r"D:\DATA\20250814EDIPLATE\trainV7_dumpTruck\models\exp8\weights\yolov7-dump-truck-20251111060914.om"
    # PT_PATH = r"yolov7-dump-truck-20251111060914.pt"
    # ONNX_PATH = r"yolov7-dump-truck-20251111060914.onnx"
    # OM_PATH = r"yolov7-dump-truck-20251111060914.om"

    # Step 1: PyTorch
    print("\n🚀 正在执行 PyTorch (.pt) 推理...")
    try:
        out_pt = infer_pytorch(PT_PATH)
        print(f"✅ .pt 输出 shape: {out_pt.shape}")
    except Exception as e:
        print(f"❌ .pt 推理失败: {e}")
        out_pt = None

    # Step 2: ONNX
    print("\n🚀 正在执行 ONNX (.onnx) 推理...")
    try:
        out_onnx = infer_onnx(ONNX_PATH)
        print(f"✅ .onnx 输出 shape: {out_onnx.shape}")
    except Exception as e:
        print(f"❌ .onnx 推理失败: {e}")
        out_onnx = None

    # Step 3: OM (仅在昇腾环境)
    print("\n🚀 正在执行 OM (.om) 推理...")
    out_om = infer_om(OM_PATH)
    if out_om is not None:
        print(f"✅ .om 输出 shape: {out_om.shape}")

    # Step 4: 对比
    print("\n" + "="*50)
    print("🔍 输出一致性对比结果:")
    print("="*50)
    if out_pt is not None and out_onnx is not None:
        compare_outputs(".pt", out_pt, ".onnx", out_onnx)
    if out_onnx is not None and out_om is not None:
        compare_outputs(".onnx", out_onnx, ".om", out_om)
    if out_pt is not None and out_om is not None:
        compare_outputs(".pt", out_pt, ".om", out_om)

    print("\n✅ 对比完成！")